package com.daxue.localclient.interceptor;

import com.daxue.localclient.utils.IDUtils;
import com.google.gson.Gson;
import io.grpc.CallOptions;
import io.grpc.Channel;
import io.grpc.ClientCall;
import io.grpc.ClientInterceptor;
import io.grpc.ForwardingClientCall;
import io.grpc.ForwardingClientCallListener;
import io.grpc.Metadata;
import io.grpc.MethodDescriptor;
import org.apache.commons.lang.StringUtils;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import org.springframework.util.StopWatch;

/**
 * @author daxue0929
 * @date 2022/08/31
 **/
public class LogGrpcInterceptor implements ClientInterceptor {

    private static final Logger log = LoggerFactory.getLogger(LogGrpcInterceptor.class);


    // 链路追踪参数traceId
    private static final Metadata.Key<String> TRACE = Metadata.Key.of("traceId", Metadata.ASCII_STRING_MARSHALLER);

    // 链路追踪参数spanId
    private static final Metadata.Key<String> SPAN = Metadata.Key.of("spanId", Metadata.ASCII_STRING_MARSHALLER);

    // 链路追踪参数parentId
    private static final Metadata.Key<String> PARENT = Metadata.Key.of("parentId", Metadata.ASCII_STRING_MARSHALLER);


    @Override
    public <ReqT, RespT> ClientCall<ReqT, RespT> interceptCall(
        MethodDescriptor<ReqT, RespT> method,
        CallOptions callOptions,
        Channel next) {
        StopWatch stopWatch = new StopWatch();
        stopWatch.start();


        log.info("发起了一次remote call 拦截住了,远程方法是：{}", method.getFullMethodName());
        try {
            return new ForwardingClientCall.SimpleForwardingClientCall<ReqT, RespT>(next.newCall(method, callOptions)) {
                @Override
                public void start(Listener<RespT> responseListener, Metadata headers) {
                    // 客户端传递链路追中数据，将数据放到headers中
                    String traceId = IDUtils.getId("trace");
                    log.info("generate traceId is: {}", traceId);
                    if (StringUtils.isNotBlank(traceId)) {
                        headers.put(TRACE, traceId);
                        headers.put(SPAN, IDUtils.getId("span"));
                        headers.put(PARENT, IDUtils.getId("parentId"));
                    }
                    // 继续下一步
                    super.start(new ForwardingClientCallListener.SimpleForwardingClientCallListener<RespT>(responseListener) {
                        @Override
                        public void onHeaders(Metadata headers) {
                            // 服务端传递回来的header
                            super.onHeaders(headers);
                            log.info("grpc call receive header：", new Gson().toJson(headers));
                        }
                    }, headers);
                }
            };
        } finally {
            stopWatch.stop();
            log.info("GRPC调用耗时： {}{}", stopWatch.getTotalTimeMillis(), "ms");

        }
    }
}
